import warnings
import time
from itertools import *
from typing import Union

import numpy as np
from scipy.stats import norm
from base_funcs import *

warnings.filterwarnings("error")


def g_1t(x, epoch):
    val = h_e * E(I(x - D_t), p_D_t) + b * E(I(D_t - x), p_D_t) + (c_e - h_m) * x
    if epoch != max_epoch:
        # last cache need to be added for f_0t
        res = alpha * E([f_0t(x - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t)
        val += res
    return val


def g_2t(x, epoch):
    val = 0
    if epoch != max_epoch:
        # last cache need to be added for f_1t
        val += alpha * E([f_1t(x - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t)
    return val


def g_3t(x, epoch):
    exp = [np.min([np.max([S_t(epoch), x]), x + K_t_i]) for K_t_i in K_t]
    val = (h_m - c_m) * x + c_m * E(exp, p_K_t)
    if epoch != max_epoch:
        val += (alpha * E([f_2t(x - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t) +
                alpha * E(
                    [f_3t(np.min([np.max([S_t(epoch), x]), x + K_t_i]) - D_t_i, epoch + 1) for K_t_i, D_t_i in K_D_t],
                    p_K_D_t))
    return val


def f_0t(v_0, epoch):
    return - c_e * v_0


def f_1t(v_1, epoch):
    return g_1t(np.max([v_1, S_2t(epoch)]), epoch)


def f_2t(v_2, epoch):
    return g_1t(np.max([v_2, S_3t(epoch)]), epoch) + g_2t(np.max([v_2, S_3t(epoch)]), epoch) - g_1t(
        np.max([v_2, S_2t(epoch)]), epoch)


def f_3t(v_3, epoch):
    return g_1t(np.max([v_3, S_4t(epoch)]), epoch) + g_2t(np.max([v_3, S_4t(epoch)]), epoch) + \
           g_3t(np.max([v_3, S_4t(epoch)]), epoch) \
           - g_1t(np.max([v_3, S_3t(epoch)]), epoch) - g_2t(np.max([v_3, S_3t(epoch)]), epoch)


def S_t(epoch):
    if not S_t_cache[epoch]["cached"]:
        if epoch == max_epoch:
            val = z[np.argmin([c_m * _z] for _z in z)]
        else:
            val = z[
                np.argmin([c_m * _z + alpha * E([f_3t(_z - D_t_i, epoch + 1) for D_t_i in D_t], p_D_t) for _z in z])]
        S_t_cache[epoch]["cached"] = True
        S_t_cache[epoch]["value"] = val
        return val
    else:
        return S_t_cache[epoch]["value"]


def S_1t(epoch):
    return S_1t_cache[epoch]["value"]


def S_2t(epoch):
    if not S_2t_cache[epoch]["cached"]:
        val = y[np.argmin([g_1t(_y, epoch) for _y in y])]
        S_2t_cache[epoch]["cached"] = True
        S_2t_cache[epoch]["value"] = val
        return val
    else:
        return S_2t_cache[epoch]["value"]


def S_3t(epoch):
    if not S_3t_cache[epoch]["cached"]:
        val = y[np.argmin([g_1t(_y, epoch) + g_2t(_y, epoch) for _y in y])]
        S_3t_cache[epoch]["cached"] = True
        S_3t_cache[epoch]["value"] = val
        return val
    else:
        return S_3t_cache[epoch]["value"]


def S_4t(epoch):
    return S_4t_cache[epoch]["value"]


def f_t(v_0, v_1, v_2, v_3, epoch):
    return f_0t(v_0, epoch) + f_1t(v_1, epoch) + f_2t(v_2, epoch) + f_3t(v_3, epoch)


def C(t):
    return c_e + np.sum([np.power(alpha, max_epoch + 1 - j) * h_e for j in range(t, max_epoch + 2)])


def init_dynamic_model(_T=3, _v_0=1, _v_1=4, _v_2=6, _v_3=8, _p=19, _b=19, _c_e=10, _c_m=2, _h_e=1, _h_m=0.2):
    global alpha, K_t, p_K_t, D_t, p_D_t, K_D_t, p_K_D_t, p, b, c_e, c_m, h_e, h_m, theta, z, y
    global epoch, max_epoch, S_t_cache, S_1t_cache, S_2t_cache, S_3t_cache, S_4t_cache
    # Variable
    alpha = 0.95
    # Kt
    mu_K_t = 6  # 均值
    cv_K_t = 0.3  # 变异系数
    sigma_K_t = cv_K_t * mu_K_t  # 标准差
    K_t = np.array([4, 5, 6, 7, 8])
    p_K_t = p_norm(K_t, mu_K_t, sigma_K_t)
    # Dt
    mu_t = 5
    sigma_t = 0.5
    D_t = np.array(list(range(0, 10)))
    p_D_t = p_norm(D_t, mu_t, sigma_t)
    # Joint Distribution of Kt, Dt
    K_D_t = [(K_t_i, D_t_i) for K_t_i, D_t_i in product(K_t, D_t)]
    # Consider as independent
    p_K_D_t = [p_K_t_i * p_D_t_i for p_K_t_i, p_D_t_i in product(p_K_t, p_D_t)]

    p = _p
    b = _b
    c_e = _c_e
    c_m = _c_m
    h_e = _h_e
    h_m = _h_m
    theta = h_m / (1 - alpha)
    print(f"c_m: {c_m}, h_e: {h_e}, h_m: {h_m}")

    max_epoch = _T - 1
    z = range(-(max_epoch + 1) * np.max(D_t), (max_epoch + 1) * np.max(K_t) + 1)
    y = range(-(max_epoch + 1) * np.max(D_t), (max_epoch + 1) * np.max(K_t) + 1)

    epoch = 1
    v_0 = _v_0
    v_1 = _v_1
    v_2 = _v_2
    v_3 = _v_3

    # S_t can be cached
    S_t_cache = {}
    S_1t_cache = {}
    S_2t_cache = {}
    S_3t_cache = {}
    S_4t_cache = {}

    # init cache
    for c in range(1, max_epoch + 2):
        S_t_cache.update({c: {"value": 0, "cached": False}})
        S_1t_cache.update({c: {"value": np.inf, "cached": True}})
        S_2t_cache.update({c: {"value": 0, "cached": False}})
        S_3t_cache.update({c: {"value": 0, "cached": False}})
        S_4t_cache.update({c: {"value": -np.inf, "cached": True}})

    calc_begin_time = time.perf_counter()
    res = f_t(v_0, v_1, v_2, v_3, epoch)
    print(res)
    calc_end_time = time.perf_counter()
    print(f"计算完成, 耗时 {calc_end_time - calc_begin_time: .2f} 秒")
    print(f"C(1): {C(1)}, C(2): {C(2)}, C(3): {C(3)}")
    return res


if __name__ == '__main__':
    print(init_dynamic_model(_T=5, _v_0=1, _v_1=8, _v_2=12, _v_3=15))
    # T = 3, 0.05s
    # T = 4, 4s
    # T = 5,
    # T = 6,
    # T = 7,
